import pandas as pd
import numpy as np
import os
import datetime
import Core.Gadget as Gadget
import Core.Config as Config
import Option.OptionGadget
import Core.WindFunctions as Wind
from SystematicFactors.General import Save_Systematic_Factor_To_Database

#
class FutureSpotBasis(object):
    def __init__(self, database, code, spot_symbol):
        self.database = database
        self.code = code
        self.spot_symbol = spot_symbol


#
class StockFutureSpotBasis(FutureSpotBasis):
    def __init__(self, database, code, spot_symbol):
        super().__init__(database, code, spot_symbol)
        self.profile_contract_count = 4 # 同一时间存在的合约数

    #
    def calc_fut_spot_basis(self, spot_symbol, fut_symbol):
        database = self.database
        #
        df_spt = database.GetDataFrame("financial_data", "index_daily_bar", filter=[("symbol", spot_symbol)], projection=["date", "close"])
        df_spt.rename(columns={"close": "spot_close"}, inplace=True)
        #
        df_fut = database.GetDataFrame("financial_data", "future_daily_bar", filter=[("symbol", fut_symbol)], projection=["date", "close", "volume"])
        df_fut.rename(columns={"close": "fut_close"}, inplace=True)
        #
        instruments = database.Find("Instruments", "Future", {"Symbol": fut_symbol})
        if len(instruments) == 0:
            print("Error, No Instrument Info", fut_symbol)
            return pd.DataFrame()

        # 确认合约的到期日
        instrument = instruments[0]
        settlement_day = instrument["datetime2"]
        # settlement_day = settlement_day.date()

        #
        df = pd.merge(df_fut, df_spt, how="left", on="date")
        df.sort_values(by="date", inplace=True)

        #
        df["day_left"] = (settlement_day - df["date"]).dt.days
        df["annualized_factor"] = 365 / df["day_left"]

        # 最后一天无法计算基差时候当NA处理,不要抛弃
        df.replace([np.inf, -np.inf], 0, inplace=True)
        # df_Fut.dropna(0, inplace=True)

        #
        df["basis"] = df["fut_close"] / df["spot_close"] - 1
        df["annual_basis"] = df["basis"] * df["annualized_factor"]

        #
        return df

    #
    def calc_weighted_basis(self, datetime1, datetime2):
        # 计算加权基差
        # weighted_field 加权字段：vol, day
        def calc_weighted_value(row, weighted_field="vol"):
            columns = row.index
            sum_value = 0
            weighted_sum_value = 0
            count = 0
            for col in columns:
                if "basis" in col:
                    x = col.split("_")
                    symbol = x[0]
                    basis = row[symbol + "_basis"]
                    value = row[symbol + "_" + weighted_field]

                    # 如果最后一天到期日，基差是0/NA
                    if np.isnan(basis):
                        continue
                    #
                    count += 1
                    sum_value += value
                    weighted_sum_value += value * basis
            #
            if count != self.profile_contract_count:
                weighted_avg_value = np.nan
            else:
                # 加权处理
                weighted_avg_value = weighted_sum_value / sum_value

            # print(row["Date"], value)
            return weighted_avg_value

        #
        trading_days_list = Gadget.GetTradingDays(self.database, datetime1, datetime2)

        #
        df = pd.DataFrame()
        dict_df_basis = {}
        symbol_list = []
        term_list = []
        for document in trading_days_list:
            day = document["date"]
            #
            # print("当前日期",day)
            # if day >= datetime.datetime(2024,5,31).date():
            #     aa = 0
            #
            terms = Option.OptionGadget.Get4Terms(day, nth=3, settlementDay=5)
            for term in terms:
                # print("   ", term)
                if term not in term_list: #
                    fut_symbol = Option.OptionGadget.term_to_symbol(self.code, term, "CFE")
                    #
                    df_basis = self.calc_fut_spot_basis(self.spot_symbol, fut_symbol)
                    df_basis.rename(columns={"annual_basis": fut_symbol + "_basis",
                                             "volume": fut_symbol + "_vol",
                                             "day_left": fut_symbol + "_day"}, inplace=True)
                    # test
                    # df_basis = df_basis[["date", fut_symbol + "_basis"]].copy()
                    df_basis = df_basis[["date", fut_symbol + "_basis", fut_symbol + "_vol", fut_symbol + "_day"]].copy()

                    #
                    if df.empty:
                        df = df_basis
                    else:
                        df = pd.merge(df, df_basis, how="outer", on="date")
                    #
                    term_list.append(term)

        #
        df.sort_values(by="date", inplace=True)
        df = df[(df["date"]>=datetime1) & (df["date"]<=datetime2)].copy()

        df["time_weighted"] = df.apply(lambda x: calc_weighted_value(x, weighted_field="day"), axis=1)
        df["vol_weighted"] = df.apply(lambda x: calc_weighted_value(x, weighted_field="vol"), axis=1)
        #
        # df.to_excel(r"C:\Users\kkwoo\Documents\Systematic_Factor\Basis/basis_test.xlsx", index=False)

        return df


def Calc_Stock_Future_Spot_Basis(database, datetime1, datetime2):

    if_datetime1 = datetime.datetime(2010, 5, 1)
    if datetime1 > if_datetime1:
        if_datetime1 = datetime1
    basis = StockFutureSpotBasis(database, code="IF", spot_symbol="000300.SH")
    df = basis.calc_weighted_basis(if_datetime1, datetime2)
    Save_Systematic_Factor_To_Database(database, df, save_name='IF_Basis_Time_Weighted', field_name="time_weighted")
    Save_Systematic_Factor_To_Database(database, df, save_name='IF_Basis_Vol_Weighted', field_name="vol_weighted")

    #
    ic_datetime1 = datetime.datetime(2015, 5, 1)
    if datetime1 > ic_datetime1:
        ic_datetime1 = datetime1
    basis = StockFutureSpotBasis(database, code="IC", spot_symbol="000905.SH")
    df = basis.calc_weighted_basis(ic_datetime1, datetime2)
    Save_Systematic_Factor_To_Database(database, df, save_name='IC_Basis_Time_Weighted', field_name="time_weighted")
    Save_Systematic_Factor_To_Database(database, df, save_name='IC_Basis_Vol_Weighted', field_name="vol_weighted")

    #
    ih_datetime1 = datetime.datetime(2015, 5, 1)
    if datetime1 > ih_datetime1:
        ih_datetime1 = datetime1
    basis = StockFutureSpotBasis(database, code="IH", spot_symbol="000016.SH")
    df = basis.calc_weighted_basis(ih_datetime1, datetime2)
    Save_Systematic_Factor_To_Database(database, df, save_name='IH_Basis_Time_Weighted', field_name="time_weighted")
    Save_Systematic_Factor_To_Database(database, df, save_name='IH_Basis_Vol_Weighted', field_name="vol_weighted")

    #
    im_datetime1 = datetime.datetime(2022, 8, 1)
    if datetime1 > im_datetime1:
        im_datetime1 = datetime1
    basis = StockFutureSpotBasis(database, code="IM", spot_symbol="000852.SH")
    df = basis.calc_weighted_basis(im_datetime1, datetime2)
    Save_Systematic_Factor_To_Database(database, df, save_name='IM_Basis_Time_Weighted', field_name="time_weighted")
    Save_Systematic_Factor_To_Database(database, df, save_name='IM_Basis_Vol_Weighted', field_name="vol_weighted")


if __name__ == '__main__':
    #
    path_filename = os.getcwd() + "\..\Config\config_local.json"
    database = Config.create_database(database_type="MySQL", config_file=path_filename, config_field="MySQL")

    #
    Wind.w.start()

    datetime1 = datetime.datetime(2024, 7, 20)
    datetime2 = datetime.datetime(2024, 11, 1)

    #
    Calc_Stock_Future_Spot_Basis(database, datetime1, datetime2)
